
// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2012 Desire NUENTSA WAKAM <desire.nuentsa_wakam@inria.fr>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#ifndef EIGEN_BROWSE_MATRICES_H
#define EIGEN_BROWSE_MATRICES_H

namespace Eigen {

enum
{
    SPD = 0x100,
    NonSymmetric = 0x0
};

/** 
 * @brief Iterator to browse matrices from a specified folder
 * 
 * This is used to load all the matrices from a folder. 
 * The matrices should be in Matrix Market format
 * It is assumed that the matrices are named as matname.mtx
 * and matname_SPD.mtx if the matrix is Symmetric and positive definite (or Hermitian)
 * The right hand side vectors are loaded as well, if they exist.
 * They should be named as matname_b.mtx. 
 * Note that the right hand side for a SPD matrix is named as matname_SPD_b.mtx
 * 
 * Sometimes a reference solution is available. In this case, it should be named as matname_x.mtx
 * 
 * Sample code
 * \code
 * 
 * \endcode
 * 
 * \tparam Scalar The scalar type 
 */
template <typename Scalar> class MatrixMarketIterator
{
    typedef typename NumTraits<Scalar>::Real RealScalar;

public:
    typedef Matrix<Scalar, Dynamic, 1> VectorType;
    typedef SparseMatrix<Scalar, ColMajor> MatrixType;

public:
    MatrixMarketIterator(const std::string& folder) : m_sym(0), m_isvalid(false), m_matIsLoaded(false), m_hasRhs(false), m_hasrefX(false), m_folder(folder)
    {
        m_folder_id = opendir(folder.c_str());
        if (m_folder_id)
            Getnextvalidmatrix();
    }

    ~MatrixMarketIterator()
    {
        if (m_folder_id)
            closedir(m_folder_id);
    }

    inline MatrixMarketIterator& operator++()
    {
        m_matIsLoaded = false;
        m_hasrefX = false;
        m_hasRhs = false;
        Getnextvalidmatrix();
        return *this;
    }
    inline operator bool() const { return m_isvalid; }

    /** Return the sparse matrix corresponding to the current file */
    inline MatrixType& matrix()
    {
        // Read the matrix
        if (m_matIsLoaded)
            return m_mat;

        std::string matrix_file = m_folder + "/" + m_matname + ".mtx";
        if (!loadMarket(m_mat, matrix_file))
        {
            std::cerr << "Warning loadMarket failed when loading \"" << matrix_file << "\"" << std::endl;
            m_matIsLoaded = false;
            return m_mat;
        }
        m_matIsLoaded = true;

        if (m_sym != NonSymmetric)
        {
            // Check whether we need to restore a full matrix:
            RealScalar diag_norm = m_mat.diagonal().norm();
            RealScalar lower_norm = m_mat.template triangularView<Lower>().norm();
            RealScalar upper_norm = m_mat.template triangularView<Upper>().norm();
            if (lower_norm > diag_norm && upper_norm == diag_norm)
            {
                // only the lower part is stored
                MatrixType tmp(m_mat);
                m_mat = tmp.template selfadjointView<Lower>();
            }
            else if (upper_norm > diag_norm && lower_norm == diag_norm)
            {
                // only the upper part is stored
                MatrixType tmp(m_mat);
                m_mat = tmp.template selfadjointView<Upper>();
            }
        }
        return m_mat;
    }

    /** Return the right hand side corresponding to the current matrix. 
     * If the rhs file is not provided, a random rhs is generated
     */
    inline VectorType& rhs()
    {
        // Get the right hand side
        if (m_hasRhs)
            return m_rhs;

        std::string rhs_file;
        rhs_file = m_folder + "/" + m_matname + "_b.mtx";  // The pattern is matname_b.mtx
        m_hasRhs = Fileexists(rhs_file);
        if (m_hasRhs)
        {
            m_rhs.resize(m_mat.cols());
            m_hasRhs = loadMarketVector(m_rhs, rhs_file);
        }
        if (!m_hasRhs)
        {
            // Generate a random right hand side
            if (!m_matIsLoaded)
                this->matrix();
            m_refX.resize(m_mat.cols());
            m_refX.setRandom();
            m_rhs = m_mat * m_refX;
            m_hasrefX = true;
            m_hasRhs = true;
        }
        return m_rhs;
    }

    /** Return a reference solution
     * If it is not provided and if the right hand side is not available
     * then refX is randomly generated such that A*refX = b 
     * where A and b are the matrix and the rhs. 
     * Note that when a rhs is provided, refX is not available 
     */
    inline VectorType& refX()
    {
        // Check if a reference solution is provided
        if (m_hasrefX)
            return m_refX;

        std::string lhs_file;
        lhs_file = m_folder + "/" + m_matname + "_x.mtx";
        m_hasrefX = Fileexists(lhs_file);
        if (m_hasrefX)
        {
            m_refX.resize(m_mat.cols());
            m_hasrefX = loadMarketVector(m_refX, lhs_file);
        }
        else
            m_refX.resize(0);
        return m_refX;
    }

    inline std::string& matname() { return m_matname; }

    inline int sym() { return m_sym; }

    bool hasRhs() { return m_hasRhs; }
    bool hasrefX() { return m_hasrefX; }
    bool isFolderValid() { return bool(m_folder_id); }

protected:
    inline bool Fileexists(std::string file)
    {
        std::ifstream file_id(file.c_str());
        if (!file_id.good())
        {
            return false;
        }
        else
        {
            file_id.close();
            return true;
        }
    }

    void Getnextvalidmatrix()
    {
        m_isvalid = false;
        // Here, we return with the next valid matrix in the folder
        while ((m_curs_id = readdir(m_folder_id)) != NULL)
        {
            m_isvalid = false;
            std::string curfile;
            curfile = m_folder + "/" + m_curs_id->d_name;
            // Discard if it is a folder
            if (m_curs_id->d_type == DT_DIR)
                continue;  //FIXME This may not be available on non BSD systems
                           //         struct stat st_buf;
                           //         stat (curfile.c_str(), &st_buf);
                           //         if (S_ISDIR(st_buf.st_mode)) continue;

            // Determine from the header if it is a matrix or a right hand side
            bool isvector, iscomplex = false;
            if (!getMarketHeader(curfile, m_sym, iscomplex, isvector))
                continue;
            if (isvector)
                continue;
            if (!iscomplex)
            {
                if (internal::is_same<Scalar, std::complex<float>>::value || internal::is_same<Scalar, std::complex<double>>::value)
                    continue;
            }
            if (iscomplex)
            {
                if (internal::is_same<Scalar, float>::value || internal::is_same<Scalar, double>::value)
                    continue;
            }

            // Get the matrix name
            std::string filename = m_curs_id->d_name;
            m_matname = filename.substr(0, filename.length() - 4);

            // Find if the matrix is SPD
            size_t found = m_matname.find("SPD");
            if ((found != std::string::npos) && (m_sym != NonSymmetric))
                m_sym = SPD;

            m_isvalid = true;
            break;
        }
    }
    int m_sym;              // Symmetry of the matrix
    MatrixType m_mat;       // Current matrix
    VectorType m_rhs;       // Current vector
    VectorType m_refX;      // The reference solution, if exists
    std::string m_matname;  // Matrix Name
    bool m_isvalid;
    bool m_matIsLoaded;  // Determine if the matrix has already been loaded from the file
    bool m_hasRhs;       // The right hand side exists
    bool m_hasrefX;      // A reference solution is provided
    std::string m_folder;
    DIR* m_folder_id;
    struct dirent* m_curs_id;
};

}  // end namespace Eigen

#endif
